from dytz.backtest import *
from dytz.params import *
from dytz.data import *
from dytz.context import *
from dytz.bar import *

import numpy as np


class DNGNode1(DNGNode):
    __params__ = {
        'k': Param(int, min_val=5, max_val=50),
        'n1': Param(int, min_val=5, max_val=50),
        'n2': Param(int, min_val=5, max_val=50),
    }

    def init_node(self):
        n = self.params['k']
        self.H = TimeSeries(n)
        self.L = TimeSeries(n)
        self.C = TimeSeries(n)
        self.O = TimeSeries(n)
        self.wl = TimeSeries(n)
        self.RSI = TimeSeries(n)

        self.collect()

    def on_bar(self, ctx: Context, bar: Bar):
        n1 = self.params['n1']
        n2 = self.params['n2']
        self.step()
        self.H.set(bar.HIGH)
        self.L.set(bar.LOW)
        self.C.set(bar.CLOSE)
        self.O.set(bar.OPEN)
        A = np.where(self.C > 0, self.C.timearray, 0).sum()
        B = np.where(self.C < 0, self.C.timearray * -1, 0).sum()
        self.RSI.set(A / (A + B) * 100)

        a = (self.H.max() - bar.CLOSE) / (self.H.max() - self.L.min()) * 100
        self.wl.set(a)
        if self.wl[0] >= n1 and self.RSI[0] > n2:
            if ctx.pos == 0:
                ctx.open_long(bar.AVGPRICE)

        if (self.wl.timearray[0] < 100 - n1 and self.wl[1] < 100 - n1) or \
                (self.RSI[0] < 100 - n2 and self.wl[0] < 100 - n1) or \
                (self.RSI[0] < 100 - n2 or self.RSI[1] < 100 - n2):
            ctx.cover_long(bar.AVGPRICE)


model = [DNGNode1(k=19, n1=9, n2=7)]

if __name__ == '__main__':
    backtest(model, '000905.SH', '20100101', '20221021', plot=True)
    # rolling_back(model, '000905.SH', '20100101', '20221021', file_name='roll')
    select_params(model, '000905.SH', '20100101', '20221021')
    # rolling_back2(model, '000905.SH', '20100101', '20221021')
